{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DeepSeek-VL2介绍"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 引言\n",
    "\n",
    "DeepSeek-VL2是一系列基于混合专家（MoE）架构的开源视觉语言模型，旨在通过三大关键升级显著提升其前身DeepSeek-VL的性能和效率。首先，DeepSeek-VL2引入了动态分块视觉编码策略，能够高效处理不同长宽比的高分辨率图像。其次，语言模型采用了多头隐变量注意力机制（Multi-head Latent Attention），通过将键值缓存压缩为潜在向量，显著提高了推理效率和吞吐量。最后，通过改进的视觉语言数据集，DeepSeek-VL2在视觉问答、光学字符识别、文档/表格/图表理解以及视觉定位等任务中展现了卓越的能力。\n",
    "\n",
    "<div style=\"display: flex; justify-content: center; align-items: center; height: [desired-container-height]px;\">\n",
    "    <img src=\"https://github.com/user-attachments/assets/fa2f662b-f084-492a-9f0f-cc5607d3ee45\" alt=\"Description\" width=\"100%\" >\n",
    "</div>\n",
    "<center>图1: DeepSeek-VL2 在视觉理解上的效果展示</center>\n",
    "\n",
    "与现有的开源密集型和基于MoE的模型相比，DeepSeek-VL2在激活参数相似或更少的情况下，实现了具有竞争力甚至最先进的性能。\n",
    "<div style=\"display: flex; justify-content: center; align-items: center; height: [desired-container-height]px;\">\n",
    "    <img src=\"https://github.com/user-attachments/assets/745b0455-a507-458f-b8e9-347f5e5c2537\" alt=\"Description\" width=\"50%\" >\n",
    "</div>\n",
    "<center>图2: 不同开源模型之间的平均性能与激活参数对比。包括MMBench v1.1、MMStar、MMMU（Val）、MathVista（TestMini）、AI2D（Test）和OCRBench</center>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 方法\n",
    "\n",
    "### 2.1 模型架构\n",
    "DeepSeek-VL2由三个核心模块组成：(1)视觉编码器，(2)视觉-语言 Adapter，以及(3)基于专家混合 (Mixture-of-Experts)的语言模型。在仅基于解码器的LLaVA架构的基础上，DeepSeek-VL2引入了两项重大改进：动态切片策略以及采用多头隐变量注意力机制的 DeepSeekMOE 语言模型。这些创新使得 DeepSeek-VL2 能够更高效地处理高分辨率视觉输入和文本数据。\n",
    "\n",
    "<div style=\"display: flex; justify-content: center; align-items: center; height: [desired-container-height]px;\">\n",
    "    <img src=\"https://github.com/user-attachments/assets/926928a3-bad2-4c5b-8f45-d0c9a7661f34\" alt=\"Description\" width=\"50%\" >\n",
    "</div>\n",
    "<center>图3: DeepSeek-VL2 架构</center>\n",
    "\n",
    "**动态切片策略**: DeepSeek-VL采用了一种混合视觉编码器，该编码器结合了SigLIP用于在384×384分辨率下进行粗粒度特征提取，以及SAM-B用于在1024×1024分辨率下进行细粒度特征提取。尽管这种融合方法能够生成适用于各种视觉-语言任务的丰富视觉表征，但它受到了固定1024×1024分辨率限制。这一局限在处理具有更大分辨率和极端长宽比的图像时尤其具有挑战性，例如信息图表视觉问答（InfographicVQA）、密集光学字符识别（OCR）以及细粒度视觉定位任务。DeepSeek-VL2将一张高分辨率图像切片，为了适应不同长宽比，首先定义了一组候选分辨率：$C_R = \\{(m \\cdot 384, n \\cdot 384)~\\mid~m \\in \\mathbb{N}, n \\in \\mathbb{N},1\\le m,n,mn\\le9\\}$, $m:n$ 表示宽高比。对于一张$(H,W)$图像，在保证宽高比不变下调整图像分辨率，计算以长边对其到候选分辨率所需要的填充的区域面积。选择面积最小的分辨率 $(m_i \\cdot 384, n_i \\cdot 384)$，然后将调整大小后的图像划分成 $m_i \\times n_i$ 个 $384 \\times 384$ 分辨率的局部图块以及一个全局缩略图块。出于计算效率和上下文长度管理的考虑，在处理多于 2 张图像时，禁用动态图块策略。\n",
    "\n",
    "<div style=\"display: flex; justify-content: center; align-items: center; height: [desired-container-height]px;\">\n",
    "    <img src=\"https://github.com/user-attachments/assets/fd5e970e-d580-4f23-81a2-4cf1f3248695\" alt=\"Description\" width=\"50%\" >\n",
    "</div>\n",
    "<center>图4: DeepSeek-VL2 中的动态切片策略</center>\n",
    "\n",
    "**视觉-语言 Adapter**: 在将图片切片后，再使用 $2 \\times 2$ pixel shuffle 操作压缩每个图像块的 token 数目， 从 $27 \\times 27$ 压缩至 $14 \\times 14=196$ tokens\n",
    "对于全局缩略图像块（$14 \\times 14$），在每一行的末尾添加14个$\\texttt{<tile\\_newline>}$标记，从而总共得到$14 \\times 15=210$ 个 tokens。\n",
    "当处理 $m_i \\times n_i$ 个局部图像块时，在每一行的局部块末尾新增$\\texttt{<tile\\_newline>}$，共新增$m_i \\cdot 14$ $\\texttt{<tile\\_newline>}$ 个 tokens，完整的 Visual Token 包含 $210 + 1 + m_i \\cdot 14 \\times (n_i \\cdot 14 + 1)$ 个视觉标记，这些 Tokens 随后使用两层多层感知器（MLP）投影到语言模型的 Embedding 空间中。\n",
    "\n",
    "**DeepSeekMoE LLM**: DeepSeek-VL2 语言模型基于DeepSeekMoE架构构建，该架构融入了多头隐变量注意力（Multi-head Latent Attention, MLA）。MLA机制通过将键值缓存压缩为潜在向量，增强了推理效率，从而提升了吞吐量。该模型还采用了混合专家（Mixture of Experts, MoE）架构，通过稀疏计算实现高效推理。在MoE训练过程中，我们为每个专家引入了一个全局偏置项，以经济高效的方式改善专家之间的负载均衡。DeepSeek-VL2推出了三种不同模型大小的变体，分别是1.0B、2.8B和4.5B。完整的架构规格如表1 所示\n",
    "<center>表1: DeepSeek-VL2 的架构规格</center>\n",
    "<div style=\"display: flex; justify-content: center; align-items: center; height: [desired-container-height]px;\">\n",
    "    <img src=\"https://github.com/user-attachments/assets/f4859490-4417-4669-9181-80f0aecd138e\" alt=\"Description\" width=\"50%\" >\n",
    "</div>\n",
    "\n",
    "现有的MoE架构可能存在知识混杂（Knowledge Hybridity）和知识冗余（Knowledge Redundancy）的问题，限制了专家的专业化。在实现思想上，DeepSeek-MoE采用两个主要策略：\n",
    "* Fine-Grained Expert Segmentation-细粒度的专家分割，通过细化FFN中间隐藏维度，维持参数数量不变的同时激活更多细粒度的专家，使得激活的专家更加灵活和适应性更强；\n",
    "\n",
    "* Shared Expert Isolation-共享专家隔离，将某些专家隔离为共享专家，始终激活，旨在捕捉和巩固不同上下文中的共同知识。\n",
    "<div style=\"display: flex; justify-content: center; align-items: center; height: [desired-container-height]px;\">\n",
    "    <img src=\"https://github.com/user-attachments/assets/e76f2103-0cd4-405b-a0be-5387a8eb6402\n",
    "    \" alt=\"Description\" width=\"50%\" >\n",
    "</div>\n",
    "<center>图5: DeepSeek-MOE 的架构</center>\n",
    "\n",
    "#### 2.1.1 细粒度的专家分割\n",
    "更可能涵盖多种类型的知识。因此，指定的专家模型会在其参数中倾向于学习大量不同类型的知识，而这些知识很难同时被利用。然而，如果每个标记（token）可以被路由到更多的专家模型中，那么不同的知识就有潜力被分别分解并在不同的专家模型中学习。在这种情况下，每个专家模型仍然可以保持高度的专业化，有助于在专家之间实现更加集中的知识分布。\n",
    "\n",
    "为了实现这一目标，同时保持专家模型参数数量和计算成本的一致性，我们以更细粒度的方式对专家进行分割。更细粒度的专家分割使得激活的专家组合更加灵活和适应性强。具体来说，在如图2(a)所示的典型混合专家（Mixture of Experts, MoE）架构之上，我们通过将前馈神经网络（Feed-Forward Network, FFN）的中间隐藏维度减少到原来的1/m倍，从而将每个专家FFN分割成m个更小的专家。由于每个专家变得更小，相应地，我们也增加了激活的专家数量到m倍，以保持相同的计算成本，如图5(b)所示。通过细粒度专家分割，MOE 层可以表示如下\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathbf{h}_{t}^{l} &= \\sum_{i=1}^{mN} \\left( g_{i,t} \\text{FFN}_{i}\\left( \\mathbf{u}_{t}^{l} \\right) \\right) + \\mathbf{u}_{t}^{l}, \\\\\n",
    "g_{i,t} &=\n",
    "\\begin{cases}\n",
    "s_{i,t}, & \\text{if } s_{i,t} \\in \\text{Topk}(\\{ s_{j, t} | 1 \\leq j \\leq mN \\}, mK), \\\\\n",
    "0, & \\text{otherwise},\n",
    "\\end{cases} \\\\\n",
    "s_{i,t} &= \\text{Softmax}_i \\left( {\\mathbf{u}_{t}^{l}}^{T} \\mathbf{e}_{i}^{l} \\right).\n",
    "\\end{aligned}\n",
    "$$\n",
    "其中，专家参数的总数等于一个标准前馈神经网络（FFN）中参数数量的 $N$ 倍。而 $mN$ 表示细粒度专家的总数。采用细粒度专家分割策略后，非零门（gate）的数量也将增加到 $mK$。为了举例说明，我们考虑$N=16$的情况。一个典型的前2名路由策略可以产生 $\\binom{16}{2}=120$ 种可能的组合。相比之下，如果每个专家被分割成4个更小的专家，那么细粒度路由策略可以产生 $\\binom{64}{8}=4,426,165,368$ 种潜在组合。组合灵活性的激增提高了实现更准确和有针对性的知识获取的可能性。\n",
    "\n",
    "#### 2.1.1 共享专家隔离\n",
    "采用传统的路由策略时，分配给不同专家的标记（tokens）可能需要一些共同的知识或信息。因此，多个专家可能会在各自的参数中收敛于获取共享知识，从而导致专家参数出现冗余。然而，如果存在专门负责在不同上下文中捕获和整合共同知识的共享专家，那么其他路由专家之间的参数冗余将得到缓解。这种冗余的缓解将有助于构建一个参数更高效、专家更专业的模型。\n",
    "\n",
    "为了实现这一目标，除了细粒度的专家分割策略外，我们还进一步隔离了 $K_{s}$ 个专家作为共享专家。无论路由器模块如何，每个标记都将被确定性地分配给这些共享专家。为了保持恒定的计算成本，其他路由专家中被激活的专家数量将减少 $K_{s}$ 个，如图5(c)所示。在集成了共享专家隔离策略后，完整的DeepSeekMoE架构中的MoE层可以表述如下：\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathbf{h}_{t}^{l} & = \\sum_{i=1}^{K_{s}} {\\operatorname{FFN}_{i}\\left( \\mathbf{u}_{t}^{l} \\right)} + \\sum_{i=K_{s} + 1}^{mN} \\left( {g_{i,t} \\operatorname{FFN}_{i}\\left( \\mathbf{u}_{t}^{l} \\right)} \\right) + \\mathbf{u}_{t}^{l}, \\\\\n",
    "g_{i,t} & = \\begin{cases} \n",
    "s_{i,t}, & s_{i,t} \\in \\operatorname{Topk} (\\{ s_{j, t} | K_{s} + 1 \\leq j \\leq mN \\}, mK - K_{s}), \\\\\n",
    "0, & \\text{otherwise}, \n",
    "\\end{cases} \\\\\n",
    "s_{i,t} & = \\operatorname{Softmax}_i \\left( {\\mathbf{u}_{t}^{l}}^{T} \\mathbf{e}_{i}^{l} \\right). \n",
    "\\end{aligned}\n",
    "$$\n",
    "最后，在DeepSeekMOE中，共享专家的数量为 $K_{s}$，路由专家的总数为$mN - K_{s}$，非零门的数量为$mK - K_{s}$。\n",
    "\n",
    "### 2.2 数据构建\n",
    "DeepSeek-VL2 从多种来源构建了一个综合性的视觉-语言数据集。训练过程分为三个阶段：（1）视觉-语言对齐（VL alignment）；（2）视觉-语言预训练（VL pretraining）；3）监督微调（Supervised Fine-Tuning）。\n",
    "\n",
    "#### 2.2.1 VL alignment 数据\n",
    "对齐阶段专注于训练多层感知机（MLP）connector，以桥接预训练的视觉编码器和大型语言模型。这一阶段使用了ShareGPT4V数据集，该数据集包含大约120万个描述和对话样本。\n",
    "\n",
    "#### 2.2.1 VL-Pretrain 数据\n",
    "VL-Pretrain 数据结合了视觉-语言数据和纯文本数据，以保持VL能力和纯文本性能之间的平衡。对于DeepSeek-VL2，我们保持了大约70%的VL数据和30%的纯文本数据的比例，后者直接来源于我们基础大型语言模型（LLM）的预训练语料库。\n",
    "\n",
    "**Image-Text 混合数据**: 数据收集始于几个开源数据集，包括WIT、WikiHow和OBELICS中的30%随机样本。这一特定的混合比例是通过使用DeepSeek-VL2-Tiny进行初步实验确定的。为了增强多语言能力，在主要以英语为主的数据集中补充了从Wanjuan中提取的中文内容。此外，DeepSeek-VL2 还开发了一个内部数据集，以扩大对一般现实世界知识的覆盖范围。\n",
    "\n",
    "**Image Caption 数据**: 图像描述是视觉语言模型（VLM）训练中的基础数据，提供了视觉信息和文本信息之间的直接对齐。因为开源数据集质量差异很大，为了解决这些质量不一致的问题，DeepSeek-VL2开发了一个全面的图像描述流程，该流程考虑了：（1）光学字符识别（OCR）提示；（2）元信息（例如位置、相机设置）；（3）原始描述作为提示。DeepSeek-VL2 使用内部Captioner，使用 类似于PixelProse的提示策略重新为图像添加描述，采用不同的指令来指导VLM生成描述。尽管Catpion整体质量有所提高，在大规模标注流程中观察到了重复问题。为了缓解这一问题，DeepSeek-VL2采用一个质量控制流程，使用DeepSeek Chat仅根据 Caption 的写作质量进行评分。\n",
    "\n",
    "**OCR 数据**: LaTex OCR 和 12M RenderedText、包括不同文档类型的大规模内部数据集\n",
    "\n",
    "**VQA 数据**: \n",
    "* DeepSeek-VL 通用的VQA数据\n",
    "* 表格、图表和文档理解数据。PubTabNet、FinTabNet 和 Docmatix\n",
    "* Web-to-code 和 plot-to-Python 生成。Websight，并遵循DeepSeek-VL的方法，使用公开的Jupyter笔记本中的Python图表。通过使用DeepSeek V2.5对Websight 部分数据增强。我们还利用DeepSeek V2.5生成的Python图表代码来减少plot-to-code 中的噪声。\n",
    "* 包括视觉提示的 QA数据：参考Vip-llava构建具有不同视觉提示（箭头、方框、圆圈和涂鸦）的数据，\n",
    "\n",
    "**Visual grounding 数据**: 基于Kosmos-2和Objects365构建 视觉定位数据，并采用以下模版构建\n",
    "* Prompt: \\texttt{Locate <|ref|><query><|/ref|> in the given image.}\n",
    "* Response: \\texttt{<|ref|><query><|/ref|><|det|>[[x1, y1, x2, y2],\\ldots]<|/det|>}\n",
    "\n",
    "**Grounded 对话数据**: 基于 Kosmos-2 构建视觉定位对话数据 \n",
    "并采用以下模版构建\n",
    "* Prompt: \\texttt{<|grounding|>Can you describe the content of the image?}\n",
    "* Response: $\\texttt{Two <|ref|>dogs<|/ref|><|det|>[[x1, y1, x2, y2],\\ldots]<|/det|> are running on the grass.}\n",
    "\n",
    "#### 2.2.1 SFT 数据\n",
    "DeepSeek-VL2 的SFT数据结合了多种开源数据集与高质量的内部QA对。\n",
    "**General visual question-answering**: 虽然VQA数据集种类繁多，但它们通常存在三大局限：（1）回答简短；（2）光学字符识别（OCR）质量不佳；（3）内容虚幻。为解决这些问题，DeepSeek-VL2 综合考虑原始问题、图像和OCR信息来重新生成回答。我们的实验表明，这种方法能产生更全面、更准确的结果。在DeepSeek-VL2的开发过程中早期版本，尤其是Tiny变体，偶尔会在中文回答中不恰当地插入英文单词。这一问题在DeepSeek-VL2 大型模型中并不存在，这表明它源于模型容量有限以及视觉-语言预训练阶段中英文数据的不平衡。为解决小型模型中的这一局限，DeepSeek-VL2团队开发了一个包含多样图像描述和单轮/多轮对话的内部中文问答数据集。该数据集有助于缓解语言混合问题。此外还创建了补充现实世界的和文化相关的视觉知识，包括动漫、网络梗、美食和艺术的内部数据集。\n",
    "\n",
    "**OCR and document understanding**: 得益于DeepSeek-VL2先进的Caption Pipeline，DeepSeek-VL2已经展现出比其他最先进的视觉语言模型（VLM）更优越的OCR能力。因此，在SFT阶段未进一步提升OCR性能，而是专注于清理现有的开源数据集，通过移除OCR质量不佳的样本。对于文档理解，DeepSeek-VL2团队从内部数据中筛选了一个多样化的文档页面子集。然后针对文档理解生成了多轮对话式问答对。\n",
    "\n",
    "**Table and chart understanding**: 通过对除Cauldron[43]（其已展现出高质量）外的所有公共数据集[14, 49]基于其原始问题重新生成回答，从而增强了基于表格的问答数据。与我们在视觉语言（VL）预训练阶段开发的OCR能力类似，我们的模型在图表理解方面也表现出色，且无需额外努力。\n",
    "\n",
    "**Textbook and academic questions**:\n",
    "我们从文档集合中构建了一个专注于教科书的内部数据集。该数据集主要强调多个学科领域的大学水平内容。\n",
    "\n",
    "**Web-to-code and plot-to-Python generation**:\n",
    "网页到代码与图表到Python代码生成。 我们扩展了内部关于网页代码和Python图表代码的数据集，这些数据集超出了预训练期间所使用的范围。对于开源数据集，我们通过重新生成答案来提高其质量。\n",
    "\n",
    "**纯文本数据**: 为了保持模型的语言能力，在SFT阶段，我们还使用了纯文本指令调整数据集\n",
    "\n",
    "### 2.3 训练方法\n",
    "#### 2.3.1 训练Pipeline\n",
    "DeepSeek-VL2通过三阶段的流程进行训练：\n",
    "* 初始阶段：在这个阶段使用第2.2.1节中详细描述的图文配对数据，训练视觉编码器和视觉-语言适配器MLP，同时保持语言模型固定。\n",
    "\n",
    "* 预训练阶段：在这一阶段，我们使用第2.2.2节描述的数据进行视觉-语言预训练。在此阶段，所有模型参数，包括视觉编码器、视觉-语言适配器和语言模型，都会解锁并同时训练。\n",
    "\n",
    "* 微调阶段：在这个阶段，我们使用第2.2.3节概述的数据进行有监督的微调，进一步优化模型性能。\n",
    "\n",
    "在预训练和微调阶段，我们强调视觉理解能力，并仅在文本标记上计算下一个标记预测损失。\n",
    "\n",
    "**视觉-语言对齐**: 基于预训练的语言模型（DeepSeekMoE 3B/16B/27B），我们的主要目标是建立视觉特征和语言特征之间的稳固连接。这种对齐使得预训练的语言模型能够有效地处理视觉输入。与之前的方法[54, 59]不同，这些方法保持预训练的视觉编码器和语言模型固定，我们调整固定分辨率的视觉编码器以适应动态高分辨率图像。在这个阶段，我们优化视觉编码器和视觉-语言适配器，同时保持语言模型冻结。\n",
    "\n",
    "**视觉-语言预训练**: 在嵌入空间中建立视觉-语言对齐之后，我们将大部分计算资源用于视觉-语言预训练。这个阶段的重点是开发跨多种任务的综合性联合视觉-语言知识。我们解锁所有参数，包括视觉编码器、视觉-语言适配器和语言模型，并同时进行训练。\n",
    "\n",
    "通过这些阶段的系统训练，DeepSeek-VL2不仅能够处理高分辨率的视觉输入，还能够在多模态任务中表现出色。这种训练方法使得模型在多样化的任务中提高了视觉和语言理解能力。\n",
    "\n",
    "**有监督微调**: 在最后阶段，我们通过有监督的微调来增强预训练模型的指令跟随能力和对话能力。利用我们内部的视觉语言SFT数据，我们优化所有参数，但仅对答案和特殊标记进行监督，同时屏蔽系统和用户提示。为了加强对话理解，我们将多模态数据与来自DeepSeek-V2的纯文本对话数据结合使用。这种方法确保了在各种视觉语言任务中具有强大的性能，包括密集图像描述、通用视觉问答（VQA）、光学字符识别（OCR）、表格/图表/文档/图形理解、视觉到代码、视觉推理、视觉定位和语言理解等。\n",
    "\n",
    "#### 3.2 超参数和基础设施\n",
    "\n",
    "DeepSeek-VL2训练的详细超参数列在表2中。我们的训练和评估是在HAI-LLM 平台上进行的，这个平台是为大型模型设计的高效轻量级平台。在我们的流水线并行策略中，视觉编码器相对于LLM块的独特计算特性带来了显著的挑战。作为模型流水线中的第一个组件，视觉编码器需要在GPU之间进行仔细的负载平衡，以防止流水线空泡并优化GPU利用率。为了解决这个问题，我们在流水线并行策略中对视觉编码器进行了细粒度的层划分。此外，我们在前向和后向过程中，在不同的数据并行等级间执行图像平铺负载平衡，以缓解动态分辨率策略导致的图像平铺数量不平衡问题。我们的训练过程还结合了张量并行和专家并行的方法，以实现最高效率。由于某些数据批次只有文本数据，而其他批次包含图像数据，我们为不同类型的数据引入了两种不同的流水线策略，并根据需要在这两种策略之间切换。DeepSeek-VL2的训练在7/10/14天内完成，使用了16/33/42个节点的集群，每个节点配备8个NVIDIA A100 GPU。\n",
    "<center>表2: DeepSeek-VL2 的训练超参数</center>\n",
    "<div style=\"display: flex; justify-content: center; align-items: center; height: [desired-container-height]px;\">\n",
    "    <img src=\"https://github.com/user-attachments/assets/9aa84dd6-5f58-41e3-87a6-6bec553f9724\" alt=\"Description\" width=\"50%\" >\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 代码解读\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 动态切片策略"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 功能:该函数 select_best_resolution 的目的是在给定的候选分辨率列表中找到最适合原始图像大小的分辨率。\n",
    "- 步骤实现：\n",
    "    - 计算缩放比例：对于每个候选分辨率，计算其相对于原始图像尺寸的缩放比例（取宽度和高度缩放比例中的最小值）。\n",
    "    - 计算缩放后的尺寸：使用上述缩放比例计算缩放后的图像宽度和高度。\n",
    "    - 计算有效分辨率：有效分辨率是缩放后的图像分辨率与原始图像分辨率中较小的一个。这是为了确保缩放后的图像不会比原始图像具有更高的分辨率。\n",
    "    - 计算浪费的分辨率：浪费的分辨率是候选分辨率的面积减去有效分辨率的面积。\n",
    "    - 选择最佳匹配：遍历所有候选分辨率，找到有效分辨率最大且浪费分辨率最小的那个作为最佳匹配。如果两个候选分辨率的有效分辨率相同，则选择浪费分辨率较小的那个。\n",
    "* 输出: 返回一个元组，代表最佳匹配的分辨率（宽度和高度）。如果没有找到任何合适的分辨率，理论上应该返回 None（尽管在当前的实现中，如果至少有一个候选分辨率，它总是会返回一个结果）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_best_resolution(image_size, candidate_resolutions):\n",
    "    original_width, original_height = image_size\n",
    "    best_fit = None\n",
    "    max_effective_resolution = 0\n",
    "    min_wasted_resolution = float(\"inf\")\n",
    "\n",
    "    for width, height in candidate_resolutions:\n",
    "        scale = min(width / original_width, height / original_height)\n",
    "        downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)\n",
    "        effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)\n",
    "        wasted_resolution = width * height - effective_resolution\n",
    "\n",
    "        if (\n",
    "            effective_resolution > max_effective_resolution\n",
    "            or effective_resolution == max_effective_resolution\n",
    "            and wasted_resolution < min_wasted_resolution\n",
    "        ):\n",
    "            max_effective_resolution = effective_resolution\n",
    "            min_wasted_resolution = wasted_resolution\n",
    "            best_fit = width, height\n",
    "\n",
    "    return best_fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 视觉-语言 Adapter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "- 方法: tokenize_with_images \n",
    "- 功能: 该函数 tokenize_with_images 的目的是将包含 $\\texttt{<image>}$ 标签的文本对话进行分词处理，并同时处理与文本对话相关联的图像。它将文本和图像转换为适合模型处理的格式，包括图像的分辨率调整、裁剪、以及将文本和图像转换为一系列 tokens。\n",
    "- 参数:\n",
    "    - conversation: 包含 $\\texttt{<image>}$ 标签的原始文本对话。\n",
    "    - images: 与文本对话中的 $\\texttt{<image>}$ 标签相对应的图像列表。\n",
    "    - bos: 布尔值，指定是否在分词结果的开头添加开始序列（Begin Of Sequence, BOS）令牌。默认为 True。\n",
    "    - eos: 布尔值，指定是否在分词结果的末尾添加结束序列（End Of Sequence, EOS）令牌。默认为 True。\n",
    "    - cropping: 布尔值，指定是否对图像进行裁剪以适应特定的分辨率。默认为 True。\n",
    "- 步骤实现：\n",
    "    - 断言检查：确保文本对话中的 $\\texttt{<image>}$ 标签数量与提供的图像数量相匹配。\n",
    "    - 文本分割：使用 $\\texttt{<image>}$ 标签将文本对话分割成多个部分。\n",
    "    - 初始化列表：用于存储处理后的图像、图像序列掩码、图像空间裁剪信息、图像令牌数量以及分词后的字符串。\n",
    "    - 遍历文本和图像：对于每个文本部分和对应的图像，执行以下操作：\n",
    "    - 文本分词：将文本部分分词，但不添加 BOS 和 EOS 令牌。\n",
    "    - 图像分辨率选择：根据裁剪标志选择最佳图像分辨率。\n",
    "    - 全局视图处理：将图像调整为固定大小（self.image_size），并填充背景色。\n",
    "    - 局部视图处理：根据最佳分辨率将图像裁剪成多个小块，并对每个小块进行处理。\n",
    "    - 记录裁剪信息：记录每个图像在宽度和高度上被裁剪成的小块数量。\n",
    "    - 添加图像token：为每个图像（全局和局部视图）生成一系列图像token，并添加到分词后的字符串中。\n",
    "    - 更新掩码和token数量：更新图像序列掩码和图像token数量列表。\n",
    "    - 处理最后一个文本部分：对最后一个文本部分进行分词处理（但不添加 BOS 和 EOS 令牌），并更新分词后的字符串和图像序列掩码。\n",
    "    - 添加 BOS 和 EOS 令牌：根据参数设置，在分词结果的开头和末尾添加 BOS 和 EOS 令牌。\n",
    "    - 断言检查：确保分词后的字符串长度与图像序列掩码的长度相匹配。\n",
    "- 输出: 返回一个元组，包含以下内容：\n",
    "    - tokenized_str: 分词后的字符串，包含文本和图像令牌。\n",
    "    - images_list: 处理后的图像列表，包括全局视图和局部视图。\n",
    "    - images_seq_mask: 图像序列掩码，用于指示哪些令牌是图像令牌。\n",
    "    - images_spatial_crop: 图像空间裁剪信息，记录每个图像在宽度和高度上的裁剪小块数量。\n",
    "    - num_image_tokens: 每个图像对应的token数量列表。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_with_images(\n",
    "    self, conversation: str, images: List[Image.Image], bos: bool = True, eos: bool = True, cropping: bool = True\n",
    "):\n",
    "    \"\"\"Tokenize text with <image> tags.\"\"\"\n",
    "    assert conversation.count(self.image_token) == len(images)\n",
    "    text_splits = conversation.split(self.image_token)\n",
    "    images_list, images_seq_mask, images_spatial_crop = [], [], []\n",
    "    num_image_tokens = []\n",
    "    tokenized_str = []\n",
    "    for text_sep, image in zip(text_splits, images):\n",
    "        \"\"\"encode text_sep\"\"\"\n",
    "        tokenized_sep = self.encode(text_sep, bos=False, eos=False)\n",
    "        tokenized_str += tokenized_sep\n",
    "        images_seq_mask += [False] * len(tokenized_sep)\n",
    "        \"\"\"select best resolution for anyres\"\"\"\n",
    "        if cropping:\n",
    "            best_width, best_height = select_best_resolution(image.size, self.candidate_resolutions)\n",
    "        else:\n",
    "            best_width, best_height = self.image_size, self.image_size\n",
    "\n",
    "        \"\"\"process the global view\"\"\"\n",
    "        global_view = ImageOps.pad(\n",
    "            image, (self.image_size, self.image_size), color=tuple(int(x * 255) for x in self.image_transform.mean)\n",
    "        )\n",
    "        images_list.append(self.image_transform(global_view))\n",
    "\n",
    "        \"\"\"process the local views\"\"\"\n",
    "        local_view = ImageOps.pad(\n",
    "            image, (best_width, best_height), color=tuple(int(x * 255) for x in self.image_transform.mean)\n",
    "        )\n",
    "\n",
    "        for i in range(0, best_height, self.image_size):\n",
    "            for j in range(0, best_width, self.image_size):\n",
    "                images_list.append(\n",
    "                    self.image_transform(local_view.crop((j, i, j + self.image_size, i + self.image_size)))\n",
    "                )\n",
    "\n",
    "        \"\"\"record height / width crop num\"\"\"\n",
    "        num_width_tiles, num_height_tiles = (best_width // self.image_size, best_height // self.image_size)\n",
    "        images_spatial_crop.append([num_width_tiles, num_height_tiles])\n",
    "\n",
    "        \"\"\"add image tokens\"\"\"\n",
    "        h = w = math.ceil(self.image_size // self.patch_size / self.downsample_ratio)\n",
    "        tokenized_image = [self.image_token_id] * h * (w + 1)\n",
    "        tokenized_image += [self.image_token_id]\n",
    "        tokenized_image += [self.image_token_id] * (num_height_tiles * h) * (num_width_tiles * w + 1)\n",
    "        tokenized_str += tokenized_image\n",
    "        images_seq_mask += [True] * len(tokenized_image)\n",
    "        num_image_tokens.append(len(tokenized_image))\n",
    "\n",
    "    \"\"\"process the last text split\"\"\"\n",
    "    tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False)\n",
    "    tokenized_str += tokenized_sep\n",
    "    images_seq_mask += [False] * len(tokenized_sep)\n",
    "\n",
    "    \"\"\"add the bos and eos tokens\"\"\"\n",
    "    if bos:\n",
    "        tokenized_str = [self.bos_id] + tokenized_str\n",
    "        images_seq_mask = [False] + images_seq_mask\n",
    "    if eos:\n",
    "        tokenized_str = tokenized_str + [self.eos_id]\n",
    "        images_seq_mask = images_seq_mask + [False]\n",
    "    assert len(tokenized_str) == len(\n",
    "        images_seq_mask\n",
    "    ), f\"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}\"\n",
    "    return (tokenized_str, images_list, images_seq_mask, images_spatial_crop, num_image_tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "### 3.3 DeepSeek-VL2 MLA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 类名:DeepseekV2Attention\n",
    "* 主要功能: 实现多头注意力机制，用于处理序列数据，支持缓存机制和不同的RoPE（Rotary Position Embedding）缩放策略。\n",
    "\n",
    "- 初始化参数 (__init__):\n",
    "    - config: DeepseekV2Config 类型的配置对象，包含模型的各种配置参数。\n",
    "    - layer_idx: 可选参数，表示当前层的索引，用于缓存机制。\n",
    "\n",
    "- 前向传播参数 (forward):\n",
    "    - hidden_states: paddle.Tensor 类型的输入张量，表示隐藏状态。\n",
    "    - attention_mask: 可选参数，paddle.Tensor 类型的注意力掩码，用于屏蔽不需要关注的位置。\n",
    "    - position_ids: 可选参数，paddle.Tensor 类型的位置编码，用于RoPE。\n",
    "    - past_key_value: 可选参数，Tuple[paddle.Tensor] 类型的缓存键值对，用于加速推理。\n",
    "    - output_attentions: 布尔类型，表示是否输出注意力权重。\n",
    "    - use_cache: 布尔类型，表示是否使用缓存机制。\n",
    "    - **kwargs: 其他可选参数。\n",
    "\n",
    "- 前向传播 (forward):\n",
    "    - 查询投影（Query Projection）：\n",
    "        - 如果 q_lora_rank 为 None，则使用 q_proj 对查询进行投影。\n",
    "        - 否则，使用基于 LoRA 的投影（q_a_proj、q_a_layernorm 和 q_b_proj）。\n",
    "\n",
    "    - 键值投影（Key-Value Projection）：\n",
    "        - 使用 kv_a_proj_with_mqa 对键和值进行投影。\n",
    "        - 将结果拆分为 LoRA 和 RoPE 组件。\n",
    "        - RoPE 应用（RoPE Application）：\n",
    "        - 计算 RoPE 的余弦和正弦值。\n",
    "        - 将 RoPE 应用于查询和键。\n",
    "    - 缓存（Caching）：\n",
    "        - 如果 use_cache 为 True，则更新缓存的键和值。\n",
    "\n",
    "    - 注意力权重（Attention Weights）：\n",
    "        - 使用缩放点积注意力计算注意力分数。\n",
    "        - 应用注意力掩码和 softmax。\n",
    "        - 输出投影（Output Projection）：\n",
    "        - 使用注意力权重和投影后的值计算注意力输出。\n",
    "        - 应用输出投影（o_proj）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepseekV2Attention(paddle.nn.Layer):\n",
    "    \"\"\"Multi-headed attention from 'Attention Is All You Need' paper\"\"\"\n",
    "\n",
    "    def __init__(self, config: DeepseekV2Config, layer_idx: Optional[int] = None):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        self.layer_idx = layer_idx\n",
    "        if layer_idx is None:\n",
    "            logger.warning_once(\n",
    "                f\"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` when creating this class.\"\n",
    "            )\n",
    "        self.attention_dropout = config.attention_dropout\n",
    "        self.hidden_size = config.hidden_size\n",
    "        self.num_heads = config.num_attention_heads\n",
    "        self.max_position_embeddings = config.max_position_embeddings\n",
    "        self.rope_theta = config.rope_theta\n",
    "        self.q_lora_rank = config.q_lora_rank\n",
    "        self.qk_rope_head_dim = config.qk_rope_head_dim\n",
    "        self.kv_lora_rank = config.kv_lora_rank\n",
    "        self.v_head_dim = config.v_head_dim\n",
    "        self.qk_nope_head_dim = config.qk_nope_head_dim\n",
    "        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim\n",
    "        self.is_causal = True\n",
    "        if self.q_lora_rank is None:\n",
    "            self.q_proj = paddle.nn.Linear(\n",
    "                in_features=self.hidden_size, out_features=self.num_heads * self.q_head_dim, bias_attr=False\n",
    "            )\n",
    "        else:\n",
    "            self.q_a_proj = paddle.nn.Linear(\n",
    "                in_features=self.hidden_size, out_features=config.q_lora_rank, bias_attr=config.attention_bias\n",
    "            )\n",
    "            self.q_a_layernorm = DeepseekV2RMSNorm(config.q_lora_rank)\n",
    "            self.q_b_proj = paddle.nn.Linear(\n",
    "                in_features=config.q_lora_rank, out_features=self.num_heads * self.q_head_dim, bias_attr=False\n",
    "            )\n",
    "        self.kv_a_proj_with_mqa = paddle.nn.Linear(\n",
    "            in_features=self.hidden_size,\n",
    "            out_features=config.kv_lora_rank + config.qk_rope_head_dim,\n",
    "            bias_attr=config.attention_bias,\n",
    "        )\n",
    "        self.kv_a_layernorm = DeepseekV2RMSNorm(config.kv_lora_rank)\n",
    "        self.kv_b_proj = paddle.nn.Linear(\n",
    "            in_features=config.kv_lora_rank,\n",
    "            out_features=self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),\n",
    "            bias_attr=False,\n",
    "        )\n",
    "        self.o_proj = paddle.nn.Linear(\n",
    "            in_features=self.num_heads * self.v_head_dim,\n",
    "            out_features=self.hidden_size,\n",
    "            bias_attr=config.attention_bias,\n",
    "        )\n",
    "        self._init_rope()\n",
    "        self.softmax_scale = self.q_head_dim**-0.5\n",
    "        if self.config.rope_scaling is not None:\n",
    "            mscale_all_dim = self.config.rope_scaling.get(\"mscale_all_dim\", 0)\n",
    "            scaling_factor = self.config.rope_scaling[\"factor\"]\n",
    "            if mscale_all_dim:\n",
    "                mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)\n",
    "                self.softmax_scale = self.softmax_scale * mscale * mscale\n",
    "\n",
    "    def _init_rope(self):\n",
    "        if self.config.rope_scaling is None:\n",
    "            self.rotary_emb = DeepseekV2RotaryEmbedding(\n",
    "                self.qk_rope_head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta\n",
    "            )\n",
    "        else:\n",
    "            scaling_type = self.config.rope_scaling[\"type\"]\n",
    "            scaling_factor = self.config.rope_scaling[\"factor\"]\n",
    "            if scaling_type == \"linear\":\n",
    "                self.rotary_emb = DeepseekV2LinearScalingRotaryEmbedding(\n",
    "                    self.qk_rope_head_dim,\n",
    "                    max_position_embeddings=self.max_position_embeddings,\n",
    "                    scaling_factor=scaling_factor,\n",
    "                    base=self.rope_theta,\n",
    "                )\n",
    "            elif scaling_type == \"dynamic\":\n",
    "                self.rotary_emb = DeepseekV2DynamicNTKScalingRotaryEmbedding(\n",
    "                    self.qk_rope_head_dim,\n",
    "                    max_position_embeddings=self.max_position_embeddings,\n",
    "                    scaling_factor=scaling_factor,\n",
    "                    base=self.rope_theta,\n",
    "                )\n",
    "            elif scaling_type == \"yarn\":\n",
    "                kwargs = {\n",
    "                    key: self.config.rope_scaling[key]\n",
    "                    for key in [\n",
    "                        \"original_max_position_embeddings\",\n",
    "                        \"beta_fast\",\n",
    "                        \"beta_slow\",\n",
    "                        \"mscale\",\n",
    "                        \"mscale_all_dim\",\n",
    "                    ]\n",
    "                    if key in self.config.rope_scaling\n",
    "                }\n",
    "\n",
    "                self.rotary_emb = DeepseekV2YarnRotaryEmbedding(\n",
    "                    self.qk_rope_head_dim,\n",
    "                    max_position_embeddings=self.max_position_embeddings,\n",
    "                    scaling_factor=scaling_factor,\n",
    "                    base=self.rope_theta,\n",
    "                    **kwargs,\n",
    "                )\n",
    "            else:\n",
    "                raise ValueError(f\"Unknown RoPE scaling type {scaling_type}\")\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        hidden_states: paddle.Tensor,\n",
    "        attention_mask: Optional[paddle.Tensor] = None,\n",
    "        position_ids: Optional[paddle.Tensor] = None,\n",
    "        past_key_value: Optional[Tuple[paddle.Tensor]] = None,\n",
    "        output_attentions: bool = False,\n",
    "        use_cache: bool = False,\n",
    "        **kwargs\n",
    "    ) -> Tuple[paddle.Tensor, Optional[paddle.Tensor], Optional[Tuple[paddle.Tensor]]]:\n",
    "\n",
    "        bsz, q_len, _ = tuple(hidden_states.shape)\n",
    "        if self.q_lora_rank is None:\n",
    "            q = self.q_proj(hidden_states)\n",
    "        else:\n",
    "            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))\n",
    "        q = q.view([bsz, q_len, self.num_heads, self.q_head_dim]).transpose(perm=[0, 2, 1, 3])\n",
    "        q_nope, q_pe = paddle.split(q, num_or_sections=[self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1)\n",
    "        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)\n",
    "        compressed_kv, k_pe = paddle.split(\n",
    "            x=compressed_kv, num_or_sections=[self.kv_lora_rank, self.qk_rope_head_dim], axis=-1\n",
    "        )\n",
    "        compressed_kv = self.kv_a_layernorm(compressed_kv)\n",
    "        k_pe = k_pe.view([bsz, q_len, 1, self.qk_rope_head_dim]).transpose(perm=[0, 2, 1, 3])\n",
    "        kv_seq_len = tuple(k_pe.shape)[-2]\n",
    "        if past_key_value is not None:\n",
    "            kv_seq_len += past_key_value[0].shape[1]\n",
    "\n",
    "        cos, sin = self.rotary_emb(q_pe, seq_len=kv_seq_len)\n",
    "        q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)\n",
    "        if use_cache and past_key_value is not None:\n",
    "            compressed_kv = compressed_kv.unsqueeze(axis=2)\n",
    "            k_pe = k_pe.transpose(perm=[0, 2, 1, 3])  # (b h l d) to (b l h d)\n",
    "            k_pe = paddle.concat([past_key_value[0], k_pe], axis=1)\n",
    "            compressed_kv = paddle.concat([past_key_value[1], compressed_kv], axis=1)\n",
    "\n",
    "            past_key_value = (k_pe, compressed_kv)\n",
    "\n",
    "            k_pe = k_pe.transpose(perm=[0, 2, 1, 3])  # go back to (b l h d)\n",
    "            compressed_kv = compressed_kv.squeeze(2)\n",
    "        elif use_cache:\n",
    "            past_key_value = (k_pe.transpose(perm=[0, 2, 1, 3]), compressed_kv.unsqueeze(axis=2))\n",
    "\n",
    "        # shit tranpose liner weight\n",
    "        kv_b_proj = self.kv_b_proj.weight.transpose(perm=[1, 0]).view([self.num_heads, -1, self.kv_lora_rank])\n",
    "        q_absorb = kv_b_proj[:, : self.qk_nope_head_dim, :]\n",
    "        out_absorb = kv_b_proj[:, self.qk_nope_head_dim :, :]\n",
    "        q_nope = paddle.matmul(x=q_nope, y=q_absorb)\n",
    "\n",
    "        perm_0 = list(range(k_pe.ndim))\n",
    "        perm_0[-1], perm_0[-2] = perm_0[-2], perm_0[-1]\n",
    "        perm_1 = list(range(compressed_kv.unsqueeze(axis=-3).ndim))\n",
    "        perm_1[-1], perm_1[-2] = perm_1[-2], perm_1[-1]\n",
    "\n",
    "        attn_weights = (\n",
    "            paddle.matmul(x=q_pe, y=k_pe.transpose(perm=perm_0))\n",
    "            + paddle.matmul(x=q_nope, y=compressed_kv.unsqueeze(axis=-3).transpose(perm=perm_1))\n",
    "        ) * self.softmax_scale\n",
    "        if tuple(attn_weights.shape) != (bsz, self.num_heads, q_len, kv_seq_len):\n",
    "            raise ValueError(\n",
    "                f\"Attention weights should be of size {bsz, self.num_heads, q_len, kv_seq_len}, but is {tuple(attn_weights.shape)}\"\n",
    "            )\n",
    "        assert attention_mask is not None\n",
    "        if attention_mask is not None:\n",
    "            if tuple(attention_mask.shape) != (bsz, 1, q_len, kv_seq_len):\n",
    "                raise ValueError(\n",
    "                    f\"Attention mask should be of size {bsz, 1, q_len, kv_seq_len}, but is {tuple(attention_mask.shape)}\"\n",
    "                )\n",
    "            attn_weights = attn_weights + attention_mask\n",
    "        attn_weights = paddle.nn.functional.softmax(x=attn_weights, axis=-1, dtype=\"float32\").to(q_pe.dtype)\n",
    "        attn_weights = paddle.nn.functional.dropout(x=attn_weights, p=self.attention_dropout, training=self.training)\n",
    "        attn_output = paddle.einsum(\"bhql,blc->bhqc\", attn_weights, compressed_kv)\n",
    "        perm_2 = list(range(out_absorb.ndim))\n",
    "        perm_2[-1], perm_2[-2] = perm_2[-2], perm_2[-1]  # tranpose the last two dims\n",
    "        attn_output = paddle.matmul(x=attn_output, y=out_absorb.transpose(perm=perm_2))\n",
    "        if tuple(attn_output.shape) != (bsz, self.num_heads, q_len, self.v_head_dim):\n",
    "            raise ValueError(\n",
    "                f\"`attn_output` should be of size {bsz, self.num_heads, q_len, self.v_head_dim}, but is {tuple(attn_output.shape)}\"\n",
    "            )\n",
    "        attn_output = attn_output.transpose(perm=[0, 2, 1, 3])\n",
    "        attn_output = attn_output.reshape([bsz, q_len, self.num_heads * self.v_head_dim])\n",
    "        attn_output = self.o_proj(attn_output)\n",
    "        if not output_attentions:\n",
    "            attn_weights = None\n",
    "        return attn_output, attn_weights, past_key_value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 DeepSeek-VL2 DeepSeekV2MOE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 类名: DeepseekV2MoE\n",
    "- 主要功能: 实现混合专家机制，通过路由机制将输入分配给多个专家网络，并将结果加权组合。\n",
    "\n",
    "- 初始化参数 (__init__):\n",
    "    - config: 配置对象，包含模型的各种参数，如专家数量、共享专家数量、中间层大小等。\n",
    "- 步骤实现\n",
    "    - 初始化 (__init__):\n",
    "        - 从配置对象中读取参数，如专家数量、共享专家数量、中间层大小等。\n",
    "        - 根据分布式环境（ep_size）分配专家网络到不同的设备上。\n",
    "        - 初始化专家网络列表（self.experts）和共享专家网络（self.shared_experts）。\n",
    "        - 初始化门控机制（self.gate）。\n",
    "\n",
    "    - 前向传播 (forward):\n",
    "        - 保存输入张量的原始形状和值（identity 和 orig_shape）。\n",
    "        - 使用门控机制（self.gate）计算路由索引（topk_idx）、路由权重（topk_weight）和辅助损失（aux_loss）。\n",
    "        - 将输入张量展平以便处理。\n",
    "        - 训练模式:\n",
    "            - 将输入张量复制多次以匹配每个专家的输入。\n",
    "            - 根据路由索引将输入分配给对应的专家网络，并计算输出。\n",
    "            - 对专家输出进行加权求和，并恢复原始形状。\n",
    "            - 添加辅助损失（AddAuxiliaryLoss.apply）。\n",
    "\n",
    "        - 推理模式:\n",
    "            - 调用 moe_infer 方法处理输入，并恢复原始形状。\n",
    "            - 如果存在共享专家网络，将其输出与专家网络的输出相加。\n",
    "            - 返回最终的输出张量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepseekV2MoE(paddle.nn.Layer):\n",
    "    \"\"\"\n",
    "    A mixed expert module containing shared experts.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        self.num_experts_per_tok = config.num_experts_per_tok\n",
    "        if hasattr(config, \"ep_size\") and config.ep_size > 1:\n",
    "            assert config.ep_size == paddle.distributed.get_world_size()\n",
    "            self.ep_size = config.ep_size\n",
    "            self.experts_per_rank = config.n_routed_experts // config.ep_size\n",
    "            self.ep_rank = paddle.distributed.get_rank()\n",
    "            self.experts = paddle.nn.LayerList(\n",
    "                sublayers=[\n",
    "                    (\n",
    "                        DeepseekV2MLP(config, intermediate_size=config.moe_intermediate_size)\n",
    "                        if i >= self.ep_rank * self.experts_per_rank and i < (self.ep_rank + 1) * self.experts_per_rank\n",
    "                        else None\n",
    "                    )\n",
    "                    for i in range(config.n_routed_experts)\n",
    "                ]\n",
    "            )\n",
    "        else:\n",
    "            self.ep_size = 1\n",
    "            self.experts_per_rank = config.n_routed_experts\n",
    "            self.ep_rank = 0\n",
    "            self.experts = paddle.nn.LayerList(\n",
    "                sublayers=[\n",
    "                    DeepseekV2MLP(config, intermediate_size=config.moe_intermediate_size)\n",
    "                    for i in range(config.n_routed_experts)\n",
    "                ]\n",
    "            )\n",
    "        self.gate = MoEGate(config)\n",
    "        if config.n_shared_experts is not None:\n",
    "            intermediate_size = config.moe_intermediate_size * config.n_shared_experts\n",
    "            self.shared_experts = DeepseekV2MLP(config=config, intermediate_size=intermediate_size)\n",
    "\n",
    "    def forward(self, hidden_states):\n",
    "        identity = hidden_states\n",
    "        orig_shape = tuple(hidden_states.shape)\n",
    "        topk_idx, topk_weight, aux_loss = self.gate(hidden_states)\n",
    "        hidden_states = hidden_states.view([-1, hidden_states.shape[-1]])\n",
    "        flat_topk_idx = topk_idx.view([-1])\n",
    "        if self.training:\n",
    "            hidden_states = hidden_states.repeat_interleave(repeats=self.num_experts_per_tok, axis=0)\n",
    "            y = paddle.empty_like(x=hidden_states)\n",
    "            for i, expert in enumerate(self.experts):\n",
    "                y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])\n",
    "            y = (y.view([*tuple(topk_weight.shape), -1]) * topk_weight.unsqueeze(axis=-1)).sum(axis=1)\n",
    "            y = y.to(hidden_states.dtype).view(orig_shape)\n",
    "            y = AddAuxiliaryLoss.apply(y, aux_loss)\n",
    "        else:\n",
    "            y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(orig_shape)\n",
    "        if self.config.n_shared_experts is not None:\n",
    "            y = y + self.shared_experts(identity)\n",
    "        return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 类名: MoEGate\n",
    "- 主要功能: 实现混合专家机制的门控逻辑，包括路由权重计算、专家选择和辅助损失计算。\n",
    "- 初始化参数 (__init__):\n",
    "    - config: 配置对象，包含模型的各种参数，如专家数量、路由缩放因子、评分函数等。\n",
    "- 步骤实现\n",
    "    - 初始化 (__init__):\n",
    "        - 从配置对象中读取参数，如专家数量、路由缩放因子、评分函数等。\n",
    "        - 初始化门控权重（self.weight）和路由策略相关参数。\n",
    "        - 如果使用 noaux_tc 路由策略，初始化专家评分校正偏置（self.e_score_correction_bias）。\n",
    "        - 调用 reset_parameters 方法初始化权重。\n",
    "\n",
    "    - 权重初始化 (reset_parameters):\n",
    "        - 使用 Kaiming 均匀分布初始化门控权重。\n",
    "\n",
    "    - 前向传播 (forward):\n",
    "        - 将输入张量展平以便处理。\n",
    "        - 使用线性变换计算路由得分（logits）。\n",
    "        - 根据评分函数（如 softmax 或 sigmoid）计算路由权重（scores）。\n",
    "        - 根据路由策略（如 greedy、group_limited_greedy、noaux_tc）选择专家并计算路由权重。\n",
    "        - 如果 top_k > 1 且 norm_topk_prob 为 True，对路由权重进行归一化。\n",
    "        - 在训练模式下，计算辅助损失（aux_loss）以优化路由机制。\n",
    "        - 返回路由索引（topk_idx）、路由权重（topk_weight）和辅助损失（aux_loss）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MoEGate(paddle.nn.Layer):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        self.top_k = config.num_experts_per_tok\n",
    "        self.n_routed_experts = config.n_routed_experts\n",
    "        self.routed_scaling_factor = config.routed_scaling_factor\n",
    "        self.scoring_func = config.scoring_func\n",
    "        self.alpha = config.aux_loss_alpha\n",
    "        self.seq_aux = config.seq_aux\n",
    "        self.topk_method = config.topk_method\n",
    "        self.n_group = config.n_group\n",
    "        self.topk_group = config.topk_group\n",
    "        self.norm_topk_prob = config.norm_topk_prob\n",
    "        self.gating_dim = config.hidden_size\n",
    "        self.weight = paddle.base.framework.EagerParamBase.from_tensor(\n",
    "            tensor=paddle.empty(shape=(self.n_routed_experts, self.gating_dim))\n",
    "        )\n",
    "        if self.topk_method == \"noaux_tc\":\n",
    "            self.e_score_correction_bias = paddle.base.framework.EagerParamBase.from_tensor(\n",
    "                tensor=paddle.empty(shape=[self.n_routed_experts])\n",
    "            )\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self) -> None:\n",
    "        init_KaimingUniform = paddle.nn.initializer.KaimingUniform(\n",
    "            negative_slope=math.sqrt(5), nonlinearity=\"leaky_relu\"\n",
    "        )\n",
    "        init_KaimingUniform(self.weight)\n",
    "\n",
    "    def forward(self, hidden_states):\n",
    "        bsz, seq_len, h = tuple(hidden_states.shape)\n",
    "        hidden_states = hidden_states.view([-1, h])\n",
    "        logits = paddle.nn.functional.linear(\n",
    "            x=hidden_states.astype(\"float32\"), weight=self.weight.astype(\"float32\").T, bias=None\n",
    "        )\n",
    "        if self.scoring_func == \"softmax\":\n",
    "            scores = paddle.nn.functional.softmax(logits, axis=-1, dtype=\"float32\")\n",
    "        elif self.scoring_func == \"sigmoid\":\n",
    "            scores = logits.sigmoid()\n",
    "        else:\n",
    "            raise NotImplementedError(f\"insupportable scoring function for MoE gating: {self.scoring_func}\")\n",
    "        if self.topk_method == \"greedy\":\n",
    "            topk_weight, topk_idx = paddle.topk(k=self.top_k, sorted=False, x=scores, axis=-1)\n",
    "        elif self.topk_method == \"group_limited_greedy\":\n",
    "            group_scores = scores.view(bsz * seq_len, self.n_group, -1).max(dim=-1).values\n",
    "\n",
    "            group_idx = paddle.topk(k=self.topk_group, sorted=False, x=group_scores, axis=-1)[1]\n",
    "            group_mask = paddle.zeros_like(x=group_scores)\n",
    "            group_mask.put_along_axis_(axis=1, indices=group_idx, values=1, broadcast=False)\n",
    "            score_mask = (\n",
    "                group_mask.unsqueeze(axis=-1)\n",
    "                .expand(shape=[bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group])\n",
    "                .reshape([bsz * seq_len, -1])\n",
    "            )\n",
    "            tmp_scores = scores.masked_fill(mask=~score_mask.astype(dtype=\"bool\"), value=0.0)\n",
    "            topk_weight, topk_idx = paddle.topk(k=self.top_k, sorted=False, x=tmp_scores, axis=-1)\n",
    "        elif self.topk_method == \"noaux_tc\":\n",
    "            assert not self.training\n",
    "            scores_for_choice = scores.view([bsz * seq_len, -1]) + self.e_score_correction_bias.unsqueeze(axis=0)\n",
    "            group_scores = scores_for_choice.view([bsz * seq_len, self.n_group, -1]).topk(k=2, axis=-1)[0].sum(axis=-1)\n",
    "\n",
    "            group_idx = paddle.topk(k=self.topk_group, sorted=False, x=group_scores, axis=-1)[1]\n",
    "            group_mask = paddle.zeros_like(x=group_scores)\n",
    "            group_mask.put_along_axis_(axis=1, indices=group_idx, values=1, broadcast=False)\n",
    "            # todo\n",
    "            score_mask = (\n",
    "                group_mask.unsqueeze(axis=-1)\n",
    "                .expand(shape=[bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group])\n",
    "                .reshape([bsz * seq_len, -1])\n",
    "            )\n",
    "            tmp_scores = scores_for_choice.masked_fill(mask=~score_mask.astype(dtype=\"bool\"), value=0.0)\n",
    "            _, topk_idx = paddle.topk(k=self.top_k, sorted=False, x=tmp_scores, axis=-1)\n",
    "            topk_weight = scores.take_along_axis(axis=1, indices=topk_idx, broadcast=False)\n",
    "        if self.top_k > 1 and self.norm_topk_prob:\n",
    "            denominator = topk_weight.sum(axis=-1, keepdim=True) + 1e-20\n",
    "            topk_weight = topk_weight / denominator * self.routed_scaling_factor\n",
    "        else:\n",
    "            topk_weight = topk_weight * self.routed_scaling_factor\n",
    "        if self.training and self.alpha > 0.0:\n",
    "            scores_for_aux = scores\n",
    "            aux_topk = self.top_k\n",
    "            topk_idx_for_aux_loss = topk_idx.view([bsz, -1])\n",
    "            if self.seq_aux:\n",
    "                scores_for_seq_aux = scores_for_aux.view([bsz, seq_len, -1])\n",
    "                ce = paddle.zeros(shape=[bsz, self.n_routed_experts])\n",
    "                ce.put_along_axis_(\n",
    "                    axis=1,\n",
    "                    indices=topk_idx_for_aux_loss,\n",
    "                    values=paddle.ones(shape=[bsz, seq_len * aux_topk]),\n",
    "                    reduce=\"add\",\n",
    "                ).divide_(y=paddle.to_tensor(seq_len * aux_topk / self.n_routed_experts))\n",
    "                aux_loss = (ce * scores_for_seq_aux.mean(axis=1)).sum(axis=1).mean() * self.alpha\n",
    "            else:\n",
    "                mask_ce = paddle.nn.functional.one_hot(\n",
    "                    num_classes=self.n_routed_experts, x=topk_idx_for_aux_loss.view([-1])\n",
    "                ).astype(\"int64\")\n",
    "                ce = mask_ce.astype(dtype=\"float32\").mean(axis=0)\n",
    "                Pi = scores_for_aux.mean(axis=0)\n",
    "                fi = ce * self.n_routed_experts\n",
    "                aux_loss = (Pi * fi).sum() * self.alpha\n",
    "        else:\n",
    "            aux_loss = None\n",
    "        return topk_idx, topk_weight, aux_loss"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
